import tensorflow as tf
import numpy as np
from sklearn.model_selection import train_test_split
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout
from tensorflow.keras.models import Sequential
from utils import read_data

data = read_data(base_path='train_224')


def create_model(input_shape):
    model = Sequential([
        Conv2D(32, (3, 3), activation='relu', input_shape=input_shape),
        MaxPooling2D(2, 2),
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Conv2D(128, (3, 3), activation='relu'),
        MaxPooling2D(2, 2),
        Flatten(),
        Dense(512, activation='relu'),
        Dropout(0.5),
        Dense(len(data), activation='softmax')  # 根据类别数量设置输出层
    ])
    return model


# 数据处理
X = []
y = []
labels = {label: idx for idx, label in enumerate(data.keys())}
print(labels)
for label, arrays in data.items():
    for array in arrays:
        X.append(array)
        y.append(labels[label])

X = np.array(X)
y = np.array(y)

# 数据切分
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 模型构建
input_shape = X_train.shape[1:]  # (512, 512, 3)
model = create_model(input_shape)

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model.fit(X_train, y_train, epochs=10, validation_data=(X_test, y_test), batch_size=16)

# 评估模型
loss, accuracy = model.evaluate(X_test, y_test)
print(f"Test accuracy: {accuracy * 100:.2f}%")

# 保存模型
model.save('cnn_model.h5')  # 替换 'path_to_my_model' 为你的保存路径
